# -*- coding: utf-8 -*-
#--------------------------------------------------------------------------
# Copyright (c) 2009 Barry Schwartz
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
# 
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.

#--------------------------------------------------------------------------

import sys
import os
import mmap
from datetime import datetime
from kt.crypto import md5_digest, encrypt_by_arcfour, encrypt_by_aes

class crypto_params:
    cipher = b'AES'
    owner_passphrase = None
    user_passphrase = None
    security_handler = 0
    encryption_algorithm = 0
    encryption_key = None
    key_length = 0
    permissions = 0
    metadata_encrypted = True

#--------------------------------------------------------------------------

cpdef encrypt_object_data_AES(object_no, generation, data):
    cdef int kl_div_8
    cdef unsigned char *p

    if isinstance(data, mmap.mmap):
        data.seek(0)
        data_string = data.read(len(data))
    else:
        data_string = data

    kl_div_8 = crypto_params.key_length // 8

    extension = b''.ljust(5)
    p = <unsigned char *> extension
    p[0] = object_no & 0xFF
    p[1] = (object_no >> 8) & 0xFF
    p[2] = (object_no >> 16) & 0xFF
    p[3] = generation & 0xFF
    p[4] = (generation >> 8) & 0xFF

    extended_key = crypto_params.encryption_key + extension + b'sAlT'
    d = md5_digest(extended_key)
    digested_key = d[:min(16, kl_div_8 + 5)]
    init_vector = os.urandom(16)

    padding_length = 16 - (len(data_string) % 16)
    padding = b''.ljust(padding_length, bytes((padding_length,)))
    assert (len(data_string) + len(padding)) % 16 == 0

    return init_vector + encrypt_by_aes(data_string + padding,
                                        digested_key, init_vector)

cpdef encrypt_object_data_ARCFOUR(object_no, generation, data):
    cdef int kl_div_8
    cdef unsigned char *p

    if isinstance(data, mmap.mmap):
        data.seek(0)
        data_string = data.read(len(data))
    else:
        data_string = data

    kl_div_8 = crypto_params.key_length // 8

    extension = b''.ljust(5)
    p = <unsigned char *> extension
    p[0] = object_no & 0xFF
    p[1] = (object_no >> 8) & 0xFF
    p[2] = (object_no >> 16) & 0xFF
    p[3] = generation & 0xFF
    p[4] = (generation >> 8) & 0xFF

    extended_key = crypto_params.encryption_key + extension
    d = md5_digest(extended_key)
    digested_key = d[:min(16, kl_div_8 + 5)]

    return encrypt_by_arcfour(data_string, digested_key)

def encrypt_object_data(object_no, generation, data):
    if crypto_params.cipher == 'AES':
        data = encrypt_object_data_AES(object_no, generation, data)
    elif crypto_params.cipher == 'ARC4':
        data = encrypt_object_data_ARCFOUR(object_no, generation, data)
    return data

#--------------------------------------------------------------------------

passphrase_padding = (b'\x28\xBF\x4E\x5E\x4E\x75\x8A\x41\x64\x00\x4E\x56\xFF\xFA' +
                      b'\x01\x08\x2E\x2E\x00\xB6\xD0\x68\x3E\x80\x2F\x0C\xA9\xFE' +
                      b'\x64\x53\x69\x7A')

def pad_passphrase(passphrase):
    return (passphrase + passphrase_padding)[:32]

def compute_encryption_key(passphrase, the_O_field, id_string):
    cdef int i
    cdef int kl_div_8
    cdef unsigned char *p
    kl_div_8 = crypto_params.key_length // 8
    perm = b'1234'
    p = <unsigned char *> perm
    p[0] = crypto_params.permissions & 0xFF
    p[1] = (crypto_params.permissions >> 8) & 0xFF
    p[2] = (crypto_params.permissions >> 16) & 0xFF
    p[3] = (crypto_params.permissions >> 24) & 0xFF
    data = (pad_passphrase(passphrase) + the_O_field + perm + id_string)
    if 4 <= crypto_params.security_handler and not crypto_params.metadata_encrypted:
        data += b'\xFF\xFF\xFF\xFF'
    d = md5_digest(data)
    if 3 <= crypto_params.security_handler:
        for i in range(0, 50):
            d = md5_digest(d[:kl_div_8])
    crypto_params.encryption_key = d[:kl_div_8]

def compute_O_field():
    cdef int i
    cdef int k
    cdef int kl_div_8
    kl_div_8 = crypto_params.key_length // 8
    if crypto_params.owner_passphrase in (None, ''):
        passphrase = crypto_params.user_passphrase
    else:
        passphrase = crypto_params.owner_passphrase
    d = md5_digest(passphrase)
    if 3 <= crypto_params.security_handler:
        for i in range(0, 50):
            d = md5_digest(d)
    arc4_key = d[:kl_div_8]
    padded_passphrase = pad_passphrase(crypto_params.user_passphrase)
    encrypted_user_passphrase = encrypt_by_arcfour(padded_passphrase, arc4_key)
    if 3 <= crypto_params.security_handler:
        key = b''.ljust(20)
        for i in range(1, 20):
            for k in range(0, kl_div_8):
                (<unsigned char *> key)[k] = (<unsigned char *> arc4_key)[k] ^ i
            encrypted_user_passphrase = encrypt_by_arcfour(encrypted_user_passphrase, key)
    return encrypted_user_passphrase

def compute_U_field(id_string):
    cdef int i
    cdef int k
    cdef int kl_div_8
    kl_div_8 = crypto_params.key_length // 8
    d = md5_digest(passphrase_padding + id_string)
    encrypted_digest = encrypt_by_arcfour(d, crypto_params.encryption_key)
    for i in range(1, 20):
        key = b''.ljust(kl_div_8)
        for k in range(0, kl_div_8):
            (<unsigned char *> key)[k] = (<unsigned char *> crypto_params.encryption_key)[k] ^ i
        encrypted_digest = encrypt_by_arcfour(encrypted_digest, key)
    result = encrypted_digest + b'0123456789012345'
    assert len(result) == 32
    return result

#--------------------------------------------------------------------------

def compute_id_string(file_location = ''):
    # Any hash that is unlikely to collide will be suitable.
    s = str(datetime.now()) + str(file_location) + str(os.urandom(16))
    if not isinstance(s, bytes):
        s = s.encode()          # Python 3.x or greater.
    return md5_digest(s)

#--------------------------------------------------------------------------
